import pymysql
from typing import List, Tuple, Any, Optional, Dict, Union

class MySQLManager:
    def __init__(self, host: str, user: str, password: str, database: Optional[str] = None, port: Optional[int] = 3306, charset: Optional[str] = 'utf8mb4'):
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.port = port
        self.charset = charset
        self.connection = None

    def connect(self) -> bool:
        try:
            self.connection = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database,
                port=self.port,
                charset=self.charset
            )
            return True
        except Exception as e:
            print(f"Failed to connect to MySQL: {e}")
            return False

# 操作
class MySQLClient:
    # 获取数据库列表
    @staticmethod
    def get_databases(mysql_manager: MySQLManager) -> List[str]:
        cursor = mysql_manager.connection.cursor()
        cursor.execute("SHOW DATABASES")
        return [row[0] for row in cursor.fetchall()]

    # 获取表列表
    @staticmethod
    def get_tables(mysql_manager: MySQLManager, database: Optional[str] = None) -> Union[List[str], Dict[str, str]]:
        cursor = mysql_manager.connection.cursor()
        # 如果没有指定数据库，使用连接时设置的数据库
        if database is None:
            if mysql_manager.database is None:
                return {"error": "没有指定数据库且连接时也未设置默认数据库"}
            database = mysql_manager.database
        cursor.execute(f"USE {database}")
        cursor.execute("SHOW TABLES")
        return [row[0] for row in cursor.fetchall()]

    # 获取指定的表结构或者全部表结构，支持表名为字符串或列表
    @staticmethod
    def get_table_structure(mysql_manager: MySQLManager, table: Optional[Union[str, List[str]]] = None, database: Optional[str] = None) -> Union[List[Dict[str, Any]], Dict[str, Dict[str, Any]], Dict[str, str]]:
        cursor = mysql_manager.connection.cursor()
        # 如果没有指定数据库，使用连接时设置的数据库
        if database is None:
            if mysql_manager.database is None:
                return {"error": "没有指定数据库且连接时也未设置默认数据库"}
            database = mysql_manager.database
        cursor.execute(f"USE {database}")
        
        # 如果table是字符串，转换为列表以便统一处理
        if isinstance(table, str):
            table = [table]
        
        if table:
            # 处理指定的表列表
            table_structures = {}
            for table_name in table:
                cursor.execute(f"DESCRIBE {table_name}")
                table_structures[table_name] = [dict(zip(("Field", "Type", "Null", "Key", "Default", "Extra"), row)) for row in cursor.fetchall()]
            
            # 如果只有一个表，直接返回该表的结构列表
            if len(table) == 1:
                return table_structures[table[0]]
            # 多个表返回字典
            return table_structures
        else:
            # 获取所有表的结构
            cursor.execute("SHOW TABLES")
            tables = [row[0] for row in cursor.fetchall()]
            table_structures = {}
            for table_name in tables:
                cursor.execute(f"DESCRIBE {table_name}")
                table_structures[table_name] = [dict(zip(("Field", "Type", "Null", "Key", "Default", "Extra"), row)) for row in cursor.fetchall()]
            return table_structures

    # 执行查询SQL语句，也支持传递数据库名，默认为连接时设置的数据库
    @staticmethod
    def execute_query(mysql_manager: MySQLManager, query: str, database: Optional[str] = None) -> Union[List[Dict[str, Any]], Dict[str, str]]:
        cursor = mysql_manager.connection.cursor()
        # 如果没有指定数据库，使用连接时设置的数据库
        if database is None:
            if mysql_manager.database is None:
                return {"error": "没有指定数据库且连接时也未设置默认数据库"}
            database = mysql_manager.database
        cursor.execute(f"USE {database}")
        cursor.execute(query)
        # 获取列名
        column_names = [desc[0] for desc in cursor.description]
        return [dict(zip(column_names, row)) for row in cursor.fetchall()]

    # 执行非查询SQL语句，也支持传递数据库名，默认为连接时设置的数据库
    @staticmethod
    def execute_non_query(mysql_manager: MySQLManager, query: str, database: Optional[str] = None, params: Optional[Tuple] = None) -> Union[int, Dict[str, str]]:
        cursor = mysql_manager.connection.cursor()
        # 如果没有指定数据库，使用连接时设置的数据库
        if database is None:
            if mysql_manager.database is None:
                return {"error": "没有指定数据库且连接时也未设置默认数据库"}
            database = mysql_manager.database
        cursor.execute(f"USE {database}")
        # 执行查询，支持参数化
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)
        mysql_manager.connection.commit()
        # 返回影响行数
        row_count = cursor.rowcount
        cursor.close()
        return row_count

